package com.zh.websocket;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.net.URI;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ZH
 * @date 23:00 2021-02-27
 */
@Component
public class MyWebsocketHandler extends TextWebSocketHandler {
    public MyWebsocketHandler() {
        System.out.println("MyWebsocketHandler constructor ... ");
    }

    static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * socket连接成功后调用
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        session.getHandshakeHeaders();
        String username = getUsername(session.getUri());
        if (isOnline(username)) {
            return;
        }

        // 用户连接成功，放入在线用户缓存
        SESSION_POOL.put(username, session);
        System.out.println("用户" + username + " 连接成功...");
    }

    private String getUsername(URI uri) {
        String query = Optional.ofNullable(uri.getQuery()).orElse("");
        String[] split = query.split("&");

        String token = Arrays.stream(split)
                .map(e -> e.split("="))
                .filter(e -> e.length == 2 && e[0].equals("token"))
                .map(e -> e[1].trim())
                .findAny().orElseThrow(() -> new IllegalArgumentException("缺少token参数"));

        return token;
    }

    /**
     * 接收消息时调用
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String username = getUsername(session.getUri());

        // 获得客户端传来的消息
        String payload = message.getPayload();
        System.out.println("server接收到[" + username + "]发送的：" + payload);
        session.sendMessage(new TextMessage("server 发送给 " + username + " 消息 " + payload + " " + LocalDateTime.now().toString()));
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        int payloadLength = message.getPayloadLength();
        System.out.println("payloadLength = " + payloadLength);
    }

    /**
     * socket断开时调用
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String username = getUsername(session.getUri());

        // 用户退出，移除缓存
        SESSION_POOL.remove(username);
        System.out.println("用户" + username + " 连接关闭...");
    }

    public boolean isOnline(String userId) {
        return SESSION_POOL.containsKey(userId);
    }

    public List<String> getAllUsers() {
        return new ArrayList<>(SESSION_POOL.keySet());
    }

    public int getCount() {
        return SESSION_POOL.size();
    }

}